import os
import pandas as pd
from volcenginesdkarkruntime import Ark

# 请确保您已将 API Key 存储在环境变量 ARK_API_KEY 中
# 初始化Ark客户端，从环境变量中读取您的API Key
client = Ark(
    # 此为默认路径，您可根据业务所在地域进行配置
    base_url="https://ark.cn-beijing.volces.com/api/v3",
    # 从环境变量中获取您的 API Key。此为默认方式，您可根据需要进行修改
    api_key='504bfd41-c285-4924-9107-74bb71846399'
    )


def ask_llm_by_row(row):
    print("----- standard request -----")
    role = "你是一个翻译专家，擅长将英文翻译成中文。"
    alpha2 = row['alpha2']
    shortname = row['shortname']
    region = row['region']
    ask_str = f"代码为{alpha2}的国家，其名称为{shortname}，所属地区为{region}，翻译成中文的国家名称是什么？"
    ask_str += f"只需要回答国家名称，不要回答其他任何内容。"
    completion = client.chat.completions.create(
    # 指定您创建的方舟推理接入点 ID，此处已帮您修改为您的推理接入点 ID
        model="doubao-1-5-thinking-pro-250415",
        messages=[
            {"role": "system", "content": role},
            {"role": "user", "content": ask_str},
        ],
    )
    print(completion.choices[0].message.content)
    return completion.choices[0].message.content


data = pd.read_csv(
    '../../wid_data/WID_countries.csv', 
    delimiter=';'
    )

# data = data.iloc[:2,:]

# dic = {}
# for index, row in data.iterrows():
#     print(index, row)
#     alpha2 = row['alpha2']
#     dic[alpha2] = ask_llm_by_row(row)

# pd.Series(dic).to_csv('en_cn_country.csv')
